from fusionmatrix.analyzers.classifier.base import Classifier, Classifiers, Option, ClassificationTask
from fusionmatrix.base.task import BaseTask, TaskExecutor
from fusionmatrix.utils import Config

from PySide6.QtWidgets import QWidget, QPushButton, QFileDialog, QMessageBox
from PySide6.QtWidgets import QHBoxLayout, QVBoxLayout, QGridLayout
from PySide6.QtWidgets import QTreeWidget, QTreeWidgetItem, QTableWidget, QTableWidgetItem
from PySide6.QtWidgets import QApplication, QTextEdit, QMenu
from PySide6.QtGui import QCursor
from PySide6.QtCore import Qt

from typing import List, Optional
from os import path


class ClassifierOptionItem(QTreeWidgetItem):
    def __init__(self, option: Option) -> None:
        super().__init__([option.name, option.explanation])
        self.option = option
        self.setChildIndicatorPolicy(QTreeWidgetItem.ChildIndicatorPolicy.DontShowIndicator)


class ClassifierItem(QTreeWidgetItem):

    def __init__(self, classifier: Classifier) -> None:
        super().__init__([classifier.name, '多分类' if classifier.multi else '单分类'])
        self.classifier = classifier
        self.item_related_columns = QTreeWidgetItem(self, ['相关列', ''])
        self.item_options = QTreeWidgetItem(self, ["可选分类项"])
        self.addChild(self.item_related_columns)
        self.addChild(self.item_options)
        self.reload()

    def reload(self) -> None:
        self.item_related_columns.setText(1, ','.join(self.classifier.related_columns))
        self.item_options.takeChildren()
        for option in self.classifier.options:
            self.item_options.addChild(ClassifierOptionItem(option))


class ClassifierEditor(QTreeWidget):

    def __init__(self, parent: 'ClassifierWidget') -> None:
        super().__init__(parent)

        self.classifier_widget = parent
        self.setColumnCount(2)
        self.setHeaderLabels(["类别", "详细说明"])
        self.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.customContextMenuRequested.connect(self.showContextMenu)

    def addClassifier(self, classifier: Classifier) -> None:
        item = ClassifierItem(classifier)
        self.addTopLevelItem(item)
        self.expandAll()
        self.resizeColumnToContents(0)

    def getClassifiers(self) -> List[Classifier]:
        classifiers = []
        for i in range(self.topLevelItemCount()):
            item = self.topLevelItem(i)
            if isinstance(item, ClassifierItem):
                classifiers.append(item.classifier)

        return classifiers

    def runClassifier(self, classifier: Classifier) -> None:
        self.classifier_widget.runClassifier(classifier)

    def showContextMenu(self, _) -> None:
        selected_items = self.selectedItems()
        if len(selected_items) == 0:
            return

        menu = QMenu()
        selected_item = selected_items[0]
        if isinstance(selected_item, ClassifierItem):
            run_action = menu.addAction("运行该分类任务")
            run_action.triggered.connect(lambda: self.runClassifier(selected_item.classifier))

        if not menu.isEmpty():
            menu.exec(QCursor.pos())


class ClassifierWidget(QWidget):

    def __init__(self, parent: Optional[QWidget] = None) -> None:
        super().__init__(parent)

        layout = QVBoxLayout()
        button_layout = QHBoxLayout()
        self.classifier_editor = ClassifierEditor(self)
        btn_load = QPushButton("载入分类任务")
        btn_save = QPushButton("保存分类任务")
        btn_run = QPushButton("运行所有分类任务")

        btn_run.clicked.connect(self.processClassificationTasks)
        btn_save.clicked.connect(self.showSaveClassificationDialog)
        btn_load.clicked.connect(self.showLoadClassificationDialog)

        button_layout.addWidget(btn_load)
        button_layout.addWidget(btn_save)
        button_layout.addWidget(btn_run)

        layout.addLayout(button_layout)
        layout.addWidget(self.classifier_editor)
        self.setLayout(layout)

    def showSaveClassificationDialog(self) -> None:
        file_name, _ = QFileDialog.getSaveFileName(self, "导出信息", "未命名分类.json", "分类配置 (*.json)")
        if file_name == "":
            return
        elif not file_name.endswith(".json"):
            file_name += ".json"

        with open(file_name, 'w', encoding='utf8') as f:
            f.write(Classifiers(classifiers=self.classifier_editor.getClassifiers()).model_dump_json(indent=4))

    def showLoadClassificationDialog(self) -> None:
        file_name, _ = QFileDialog.getOpenFileName(self, "载入信息", "", "分类配置 (*.json)")
        if file_name == "":
            return

        with open(file_name, 'r', encoding='utf8') as f:
            self.classifier_editor.clear()
            classifiers = Classifiers.model_validate_json(f.read()).classifiers
            for classifier in classifiers:
                self.classifier_editor.addClassifier(classifier)

    def runClassifier(self, classifier: Classifier) -> None:
        self.runClassifiers([classifier])

    def runClassifiers(self, classifiers: List[Classifier]) -> None:
        from fusionmatrix.windows.main import MainWindow

        table = MainWindow.getActiveTable()
        for classifier in classifiers:
            table.createColumnHeader(classifier.name)

        model = Config.load_default_config().chat_model

        for classifier in classifiers:
            if not table.hasColumnHeader(classifier.name):
                table.createColumnHeader(classifier.name)

        tasks: List[BaseTask] = []
        for i in range(table.rowCount()):
            row = table.getRowContent(i)
            for classifier in classifiers:
                task = ClassificationTask(
                    row=row, model=model,
                    classifier=classifier,
                    target=(table.name, i, table.getColumnIndex(classifier.name))
                )
                tasks.append(task)

        TaskExecutor.submitBatchTasks(tasks)

    def processClassificationTasks(self) -> None:
        classifiers = self.classifier_editor.getClassifiers()
        self.runClassifiers(classifiers)
